/**
 * 
 */
package edu.umd.clip.lm.model.decoding;

import java.util.*;

import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.data.Context;
import edu.umd.clip.lm.storage.RequestBundle;
import edu.umd.clip.lm.storage.AbstractProbTreeStorage.Key;
import edu.umd.clip.lm.util.tree.BinaryTree;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class DecodingContext {
	private final IdentityHashMap<TrellisNodeItem, ItemData> perItemData;
	private final TreeTrellisNode currentNode;
	private final Context overtContext;
	private final ForestModel forest;
	private final FuturePredictions<TrellisNodeItem> predictions;

	private final RequestBundle<Key,OnDiskCompactProbTree> currentRequests;
	private final double prevNodeTotalProb;

	public DecodingContext(TreeTrellisNode currentNode, ForestModel forest) {
		this.currentNode = currentNode;
		this.forest = forest;
		
		this.prevNodeTotalProb = currentNode.getPreviousNode() == null ? 1.0 : currentNode.getPreviousNode().getTotalProb();
		
		this.overtContext = new Context(forest.getOvertOrder());
		
		TreeTrellisNode node = currentNode.getPreviousNode();
		for(int i=overtContext.data.length-1; i>=0 && node != null; --i, node = node.getPreviousNode()) {
			overtContext.data[i] = node.getCurrentWord();
		}

		this.perItemData = new IdentityHashMap<TrellisNodeItem, ItemData>();
		this.predictions = new FuturePredictions<TrellisNodeItem>();
		this.currentRequests = new RequestBundle<Key,OnDiskCompactProbTree>(forest.getDecoder().getStorage());
	}
	
	public final static class ItemData implements Cloneable {
		public BinaryTree<HistoryTreePayload> decodingPosition[];
		
		/**
		 * @param decodingPosition
		 */
		public ItemData(BinaryTree<HistoryTreePayload>[] decodingPosition) {
			this.decodingPosition = decodingPosition;
		}

		public BinaryTree<HistoryTreePayload> getPosition(int lmid) {
			return decodingPosition[lmid];
		}
		
		public void setPosition(int lmid, BinaryTree<HistoryTreePayload> node) {
			decodingPosition[lmid] = node;
		}

		@Override
		public Object clone() {
			return new ItemData(decodingPosition.clone());
		}
	}
	
	
	public ItemData getItemData(TrellisNodeItem item) {
		return perItemData.get(item);
	}
	
	public void setItemData(TrellisNodeItem item, ItemData data) {
		perItemData.put(item, data);
	}

	public void addPrediction(BinaryTree<HistoryTreePayload>[] nodes, TrellisNodeItem prevItem) {
		predictions.addPrediction(nodes, prevItem);
	}

	public FuturePredictions<TrellisNodeItem> getPredictions() {
		return predictions;
	}
	
	public TreeTrellisNode getCurrentNode() {
		return currentNode;
	}

	public Context getOvertContext() {
		return overtContext;
	}

	public double getPrevNodeTotalProb() {
		return prevNodeTotalProb;
	}
	
	public TreeTrellisNode getNodeByOffset(int offset) {
		TreeTrellisNode node = currentNode;
		while(offset > 0 && node != null) {
			node = node.getPreviousNode();
			--offset;
		}
		return node;
	}

	public RequestBundle<Key, OnDiskCompactProbTree> getCurrentRequests() {
		return currentRequests;
	}

	public long getFutureWord() {
		return currentNode.getCurrentWord();
	}
}
